import random
from typing import List, Tuple

class PromptConstructor:
    def __init__(self, all_cn_diseases: List[str], stage_map: dict, stage_group: dict):
        self.all_cn_diseases = all_cn_diseases
        self.stage_map = stage_map
        self.stage_group = stage_group

    def construct_prompts(self, true_diag: List[str]) -> List[Tuple[str, str, str]]:
        """
        true_diag: 可能为单个或多个疾病的list
        返回：(问题, 答案, 题型tag)元组列表，题型tag为'judge'/'diagnosis'/'compare'
        """
        prompts = []

        # 1. 判断题
        judge_templates = [
            "图像中显示的是{}吗？请严格使用“是”或“否”回答。",
            "患者是否患有{}？请直接回答是或否。",
            "该眼底照片能诊断为{}吗？请明确回答是或否。",
            "请判断图像所示是否为{}，回答是或否。",
            "是否可以从该图像诊断出{}？请明确回答是或否。"
        ]

        for diag in true_diag:
            tpl = random.choice(judge_templates)
            question = tpl.format(diag)
            answer = "是。"
            prompts.append((question, answer, "judge"))
        wrong_candidates = [d for d in self.all_cn_diseases if d not in true_diag]
        wrong_candidates = random.sample(wrong_candidates, min(2, len(wrong_candidates)))
        for diag in wrong_candidates:
            tpl = random.choice(judge_templates)
            question = tpl.format(diag)
            answer = "否。"
            prompts.append((question, answer, "judge"))

        # 2. 诊断题
        diagnosis_templates = [
            "请根据图像直接给出诊断结论。",
            "该患者的眼底情况如何？请直接诊断。",
            "根据图像，你认为患者的诊断结果是什么？",
            "请直接指出本图像的疾病诊断结果。",
            "分析眼底照片后，给出你的诊断结论。"
        ]
        answer_templates = [
            "本张图片的诊断结果是{}。",
            "该病例诊断为{}。",
            "经判断，本例患者为{}。",
            "此图像显示患者诊断为{}。",
            "眼底图片显示的疾病诊断为{}。"
        ]
        ans_str = "，".join(true_diag)
        q_diag = random.choice(diagnosis_templates)
        a_diag = random.choice(answer_templates).format(ans_str)
        prompts.append((q_diag, a_diag, "diagnosis"))

        # 3. 比较题
        compare_templates = [
            "图像更可能属于以下哪种疾病：{}或{}？请直接给出疾病名称。",
            "患者的眼底情况更符合{}还是{}？请直接回答疾病名称。",
            "根据图像，{}和{}中，哪种疾病的可能性更高？请直接回答疾病名称。",
            "从图像来看，患者的诊断更倾向于{}还是{}？直接给出诊断名称即可。",
            "请判断该图像更接近以下哪种疾病：{}或者{}？直接回答疾病名称即可。"
        ]
        compare_answer_templates = [
            "{}",
            "我认为更可能是{}。",
            "图像更接近{}。",
            "该患者诊断为{}。",
            "{}的可能性更高。",
            "诊断为{}。",
            "我的判断是{}。"
        ]
        for diag in true_diag:
            curr_wrong_candidates = [d for d in self.all_cn_diseases if d not in true_diag]
            if len(curr_wrong_candidates) >= 2:
                wrongs = random.sample(curr_wrong_candidates, 2)
            elif curr_wrong_candidates:
                wrongs = curr_wrong_candidates
            else:
                wrongs = []
            for wrong_diag in wrongs:
                tpl = random.choice(compare_templates)
                if random.random() < 0.5:
                    question = tpl.format(diag, wrong_diag)
                else:
                    question = tpl.format(wrong_diag, diag)
                ans_tpl = random.choice(compare_answer_templates)
                answer = ans_tpl.format(diag)
                prompts.append((question, answer, "compare"))

        return prompts
